﻿using System;
using System.Collections.Generic;
using Allegro.Framework;

namespace Allegro.Mathlib
{
    /// <summary>
    /// 
    /// </summary>
    public class LinearEquationSet
    {
        const double TINY = 1.0e-20;

        private int _n;
        private double[,] _a;
        private double[] _vv; // vv stores the implicit scaling of each row
        private double _d;
        private double[,] _inverseA;

        /// <summary>
        /// Initialize an equation set for the vector equation A x = b
        /// </summary>
        /// <param name="a"></param>
        public LinearEquationSet(double[,] a)
        {
            // Verify that this is a square matrix
            int n = a.GetLength(0);
            if (n != a.GetLength(1))
            {
                string msg = "Input must be a square matrix";
                ErrorHandler.Error(msg);
            }
            _n = n;
            _a = a;
            _vv = new double[n];
            _d = 1.0;
            Rows = new double[n][];
            RowPermutations = new int[n];
            for (int i = 0; i < n; i++)
            {
                RowPermutations[i] = i;
                Rows[i] = new double[n];
                for (int j = 0; j < n; j++)
                {
                    Rows[i][j] = a[i, j];
                }
            }
            // Compute the LU decomposition
            _LUDecompose();
        }

        /// <summary>
        /// Number of equations
        /// </summary>
        public int N
        {
            get { return _n; }
        }
        /// <summary>
        /// Matrix A in the vector equation A x = b
        /// </summary>
        public double[,] A
        {
            get { return _a; }
            set
            {
                if (_a.GetLength(0) == _n && _a.GetLength(1) == _n)
                {
                    _a = value;
                    _Reinitialize(value);
                }
            }
        }
        /// <summary>
        /// Rows
        /// </summary>
        public double[][] Rows;

        /// <summary>
        /// Row permutations after LU decomposition
        /// </summary>
        public int[] RowPermutations;

        /// <summary>
        /// Solve a set of linear equations defined by the matrix equation
        /// a x = b
        /// </summary>
        /// <param name="A">Matrix A</param>
        /// <param name="b">RHS vector b</param>
        /// <param name="x">Result vector</param>
        public void Solve(double[] b, ref double[] x)
        {
            // Check dimension
            if (b.Length != _n)
                ErrorHandler.Error("RHS of equation has wrong dimension");

            for (int i = 0; i < _n; i++)
                x[i] = b[i];

            // Do back substitution, using the already available LU decomposition
            _LUBackSubstitution(ref x);
        }

        /// <summary>
        /// Return the inverse of the input matrix
        /// </summary>
        public double[,] InverseMatrix
        {
            get
            {
                if (_inverseA == null)
                {
                    _inverseA = new double[_n, _n];
                    _Invert();
                }
                return _inverseA;
            }
        }

        // Private members
        void _Reinitialize(double[,] a)
        {
            _d = 1.0;
            for (int i = 0; i < _n; i++)
            {
                RowPermutations[i] = i;
                //Rows[i] = new double[n];
                for (int j = 0; j < _n; j++)
                {
                    Rows[i][j] = a[i, j];
                }
            }
            // Compute the LU decomposition
            _LUDecompose();
        }
        void _Invert()
        {
            double[] col = new double[_n];

            // Find inverse by columns
            for (int j = 0; j < _n; j++)
            {
                for (int i = 0; i < _n; i++)
                    col[i] = 0.0;
                col[j] = 1.0;

                _LUBackSubstitution(ref col);
                for (int i = 0; i < _n; i++)
                    _inverseA[i, j] = col[i];
            }
        }
        /// <summary>
        /// Provide an LU decomposition of the equation set
        /// Before the 
        /// </summary>
        void _LUDecompose()
        {
            int i, imax = 0, j, k;
            double big, dum, sum, temp;

            // Loop over rows to get the implicit scaling information
            for (i = 0; i < _n; i++)
            {
                big = 0.0;
                for (j = 0; j < _n; j++)
                {
                    temp = Math.Abs(Rows[i][j]);
                    if (temp > big)
                        big = temp;
                }
                if (big == 0.0) // No nonzero largest element
                    ErrorHandler.Error("LinearEquationSet.LUDecompose: Singular matrix");
                _vv[i] = 1.0 / big;
            }
            // Loop over columns - Crout's method
            for (j = 0; j < _n; j++)
            {
                // i < j
                for (i = 0; i < j; i++)
                {
                    sum = Rows[i][j];
                    for (k = 0; k < i; k++)
                        sum -= Rows[i][k] * Rows[k][j];
                    Rows[i][j] = sum;
                }
                // Initialize for the search for the largest pivot element
                big = 0.0;
                // i = j
                for (i = j; i < _n; i++)
                {
                    sum = Rows[i][j];
                    for (k = 0; k < j; k++)
                        sum -= Rows[i][k] * Rows[k][j];
                    Rows[i][j] = sum;

                    dum = _vv[i] * Math.Abs(sum);
                    if (dum >= big)
                    {
                        // Is the figure of merit for the pivot better than the best so far?
                        big = dum;
                        imax = i;
                    }
                }
                // Do we need to interchange rows?
                if (j != imax)
                {
                    // Yes, do so... Switch rows j and imax
                    double[] row = Rows[imax];
                    Rows[imax] = Rows[j];
                    Rows[j] = row;
                    _d = -_d; // Change the parity of _d
                    _vv[imax] = _vv[j]; // Interchange the scale factor
                }
                RowPermutations[j] = imax;
                if (Rows[j][j] == 0.0)
                    Rows[j][j] = TINY;
                // If the pivot element is zero, the matrix is singular (at least to the precision of the
                // algorithm). For some applications on singular matrices, it is desirable to substistute
                // TINY for zero

                // Finally, divide by the pivot element
                if (j != _n)
                {
                    dum = 1.0 / Rows[j][j];
                    for (i = j + 1; i < _n; i++)
                        Rows[i][j] *= dum;
                }
            } // End column loop
        }

        void _LUBackSubstitution(ref double[] bx)
        {
            // Sovles the set of _n linear equations A x = b. Here A[1..n] is input as the LU decomposition
            // of the original LHS matrix A. b inputs as the RHS vector b and returns as the solution vector x
            // This function takes into account the possibility that bx will begin with many zero elements,
            // so it is efficient for use in matrix inversion
            int i, ii = -1, ip, j;
            double sum;

            for (i = 0; i < _n; i++)
            {
                ip = RowPermutations[i];
                sum = bx[ip];
                bx[ip] = bx[i];
                // When ii is set to a positive value, it will become the index of the first nonvanishing element
                // of bx. We now do the forward substitution. The only new wrinkle is to unscramble the permutation
                // as we go
                if (ii >= 0)
                {
                    for (j = ii; j <= i - 1; j++)
                        sum -= Rows[i][j] * bx[j];
                }
                else if (sum != 0.0)
                    ii = i; // A nonzero element was encountered, so from now on we will have to do the sums in the loop above
                bx[i] = sum;
            }
            // Now we do the backsubstitution
            for (i = _n - 1; i >= 0; i--)
            {
                sum = bx[i];
                for (j = i + 1; j < _n; j++)
                    sum -= Rows[i][j] * bx[j];
                // Store the solution vector
                bx[i] = sum / Rows[i][i];
            }
        }
    }

}
